﻿#pragma once

#include "gopt_base.h"
#include "lbfgs.hpp"
#include <iomanip>
#include <iostream>


namespace GOptSolver {

struct LBFGSState {
  double iteration = -1; // iter
  double step = 0;       // line search step
  double ls = 0;         // The number of evaluations called for this iteration.
};

using print_lbfgs_progress_func_type =
    std::function<int(LBFGSState &state, const double fx,
                      const Eigen::VectorXd &x, const Eigen::VectorXd &g)>;

inline int default_lbfgs_print_progress(LBFGSState &state, const double fx,
                                        const Eigen::VectorXd &x,
                                        const Eigen::VectorXd &g) {
  std::cout << std::setprecision(4)
            << "=============LBFGS===================" << std::endl
            << "Iteration: " << state.iteration << std::endl
            << "Function Value: " << fx << std::endl
            << "Gradient Inf Norm: " << g.cwiseAbs().maxCoeff() << std::endl
            << "Variables: " << std::endl
            << x.transpose() << std::endl;
  return 0;
}

class LBFGSSolver : public BaseSolver {
private:
  evaluate_func_type objective_;
  print_lbfgs_progress_func_type print_progress_ = default_lbfgs_print_progress;
  void *data_;
  LBFGSState state_;

  static double default_lbfgs_evaluate(void *self, const Eigen::VectorXd &x,
                                       Eigen::VectorXd &grad) {
    auto ptr = static_cast<LBFGSSolver *>(self);
    return ptr->objective_(ptr->data_, x, grad);
  }

  static int default_local_monitor(void *self, const Eigen::VectorXd &x,
                                   const Eigen::VectorXd &g, const double fx,
                                   const double step, const int k,
                                   const int ls) {
    auto ptr = static_cast<LBFGSSolver *>(self);
    ptr->state_.ls = ls;
    ptr->state_.step = step;
    ptr->state_.iteration = k;
    if (ptr->print_progress_) {
      return ptr->print_progress_(ptr->state_, fx, x, g);
    }
    return 0;
  }

public:
  lbfgs::lbfgs_parameter_t param;
  LBFGSSolver() = default;
  ~LBFGSSolver() = default;
  void set_objective(const evaluate_func_type &objective,
                     void *instance) override {
    objective_ = objective;
    data_ = instance;
  }
  int minimize(Eigen::VectorXd &x, double &minf) override {
    if (!objective_) {
      return lbfgs::LBFGSERR_INVALID_FUNCVAL;
    }
    return lbfgs::lbfgs_optimize(x, minf, default_lbfgs_evaluate, nullptr,
                                 default_local_monitor, this, this->param);
  }
  const LBFGSState &get_state() const { return state_; }
  void set_progress(const print_lbfgs_progress_func_type &print_progress) {
    print_progress_ = print_progress;
  }
};

} // namespace GOptSolver
